import base64

from PIL import Image
from flask import (Flask, jsonify, request, current_app)
from GAN.visualize import *
from fastai.vision import *
torch.backends.cudnn.benchmark = True
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
app = Flask(__name__)

colorizer = get_image_colorizer()

@app.route('/api/input', methods=['POST'])
def input_image():
    data = {"success": 500}
    f = request.files['files']
    render_factor = int(request.form['render_factor'])

    image = f.read()
    img = PIL.Image.open(BytesIO(image)).convert('RGB')

    # render_factor = 35
    image = colorizer.get_transformed_image(orig_image=img, render_factor=render_factor)
    IO = BytesIO()
    image.save(IO, format='jpeg')
    # 保存在本地
    # result_path="result_images/"+request.form['render_factor']+".jpg"
    # image.save(result_path)
    # 从内存中取出bytes类型的图片
    IOImg = IO.getvalue()

    data["success"] = 200
    data.update(msg="上传成功")
    # 进行base64编码
    data.update(img_base64=base64.b64encode(IOImg).decode())
    return jsonify(data)


@app.route('/test')
def hello():
    return 'hello world'


@app.route('/favicon.ico')
def favicon():
    return current_app.send_static_file('favicon.ico')


if __name__ == "__main__":
    # app.debug = True
    app.run(host='0.0.0.0', port=5000)
